import asyncio
import json
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import httpx
from datetime import datetime, timezone, timedelta
from pydantic import AnyUrl, BaseModel

from mcp import ClientSession
from mcp.client.auth import OAuthClientProvider, TokenStorage
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import OAuthClientInformationFull, OAuthClientMetadata, OAuthToken


SERVER_URL = "https://mcp.notion.com/mcp"


class StoredToken(BaseModel):
    token_payload: OAuthToken
    expires_at: datetime | None


class FileTokenStorage(TokenStorage):
    """File-backed token storage."""

    def __init__(self, tokens_file: Path = Path("tokens.json"), client_info_file: Path = Path("client_info.json")):
        self.tokens_file = tokens_file
        self.client_info_file = client_info_file

    async def get_tokens(self) -> OAuthToken | None:
        if not self.tokens_file.exists():
            return None

        try:
            stored_token = StoredToken.model_validate_json(self.tokens_file.read_text())
        except (OSError, json.JSONDecodeError):
            return None

        if stored_token.expires_at is not None:
            now = datetime.now(timezone.utc)

            if stored_token.token_payload.expires_in is not None:
                remaining = stored_token.expires_at - now
                stored_token.token_payload.expires_in = max(0, int(remaining.total_seconds()))
        
        return stored_token.token_payload

    async def set_tokens(self, tokens: OAuthToken) -> None:
        expires_at: datetime | None = None

        if tokens.expires_in is not None:
            expires_at = datetime.now(timezone.utc) + timedelta(seconds=tokens.expires_in)

        stored_token = StoredToken(token_payload=tokens, expires_at=expires_at)

        self.tokens_file.write_text(stored_token.model_dump_json())

    async def get_client_info(self) -> OAuthClientInformationFull | None:
        if not self.client_info_file.exists():
            return None

        try:
            return OAuthClientInformationFull.model_validate_json(self.client_info_file.read_text())
        except (OSError, json.JSONDecodeError):
            return None

    async def set_client_info(self, client_info: OAuthClientInformationFull) -> None:
        self.client_info_file.write_text(client_info.model_dump_json())


async def handle_redirect(auth_url: str) -> None:
    print(f"Visit: {auth_url}")


async def handle_callback() -> tuple[str, str | None]:
    callback_url = input("Paste callback URL: ")
    params = parse_qs(urlparse(callback_url).query)
    return params["code"][0], params.get("state", [None])[0]


async def main():
    """Run the OAuth client example."""
    oauth_auth = OAuthClientProvider(
        server_url=SERVER_URL,
        client_metadata=OAuthClientMetadata(
            client_name="Example MCP Client",
            redirect_uris=[AnyUrl("http://localhost:3000/callback")],
            grant_types=["authorization_code", "refresh_token"],
            response_types=["code"],
            scope="user",
        ),
        storage=FileTokenStorage(),
        redirect_handler=handle_redirect,
        callback_handler=handle_callback,
    )

    async with httpx.AsyncClient(auth=oauth_auth, follow_redirects=True) as custom_client:
        async with streamable_http_client(SERVER_URL, http_client=custom_client) as (read, write, _):
            async with ClientSession(read, write) as session:
                await session.initialize()

                tools = await session.list_tools()
                print(f"Available tools: {[tool.name for tool in tools.tools]}")

                resources = await session.list_resources()
                print(f"Available resources: {[r.uri for r in resources.resources]}")


def run():
    asyncio.run(main())


if __name__ == "__main__":
    run()